package main

import (
	"bufio"
	"io"
	"strings"
	"sync"
	"unicode/utf8"
)

// Trie 前缀树结构
type Trie struct {
	Root *TrieNode
	mu   sync.RWMutex
}

// TrieNode 前缀树节点
type TrieNode struct {
	Children sync.Map // 修改为 sync.Map
	IsEnd    bool
}

// NewTrie 创建新的前缀树
func NewTrie() *Trie {
	return &Trie{
		Root: &TrieNode{
			Children: sync.Map{}, // 初始化为 sync.Map
			IsEnd:    false,
		},
	}
}

// findLongestWordInTrie 查找最长的违禁词
func findLongestWordInTrie(subText string, t *Trie) (string, int) {
	node := t.Root
	var word []rune
	for _, char := range subText {
		childNode, ok := node.Children.Load(char) // 使用 Load 方法获取子节点
		if ok {
			word = append(word, char)
			node = childNode.(*TrieNode) // 类型断言
			if node.IsEnd {
				return string(word), len(word)
			}
		} else {
			break
		}
	}
	return "", 0
}

// Get 获取文本中包含的所有违禁词
func (t *Trie) Get(text string) []string {
	t.mu.RLock()
	defer t.mu.RUnlock()
	var result = make([]string, 0)
	uniqueMap := make(map[string]bool)

	text = strings.ReplaceAll(text, " ", "")
	textRunes := []rune(text)

	for i := 0; i < len(textRunes); {
		subText := string(textRunes[i:])
		if word, size := findLongestWordInTrie(subText, t); word != "" {
			if !uniqueMap[word] {
				uniqueMap[word] = true
				result = append(result, word)
			}
			i += size
		} else {
			i++
		}
	}
	return result
}

// Is 检测文本中是否包含任何违禁词
func (t *Trie) Is(text string) bool {
	t.mu.RLock()
	defer t.mu.RUnlock()
	text = strings.ReplaceAll(text, " ", "")
	for i := 0; i < len(text); {
		// 检测从当前位置开始的子串
		if _, size := findLongestWordInTrie(text[i:], t); size > 0 {
			return true
		}
		i += utf8.RuneCountInString(text[i:])
	}
	return false
}

// Add 插入单词到前缀树
func (t *Trie) Add(word ...string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, w := range word {
		node := t.Root
		for _, char := range w {
			// 使用 LoadOrStore 方法，如果不存在则创建并存储
			actualNode, _ := node.Children.LoadOrStore(char, &TrieNode{
				Children: sync.Map{},
				IsEnd:    false,
			})
			node = actualNode.(*TrieNode) // 类型断言
		}
		node.IsEnd = true
	}
}

// Remove 从前缀树中移除单词
func (t *Trie) Remove(words ...string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, word := range words {
		node := t.Root
		for _, char := range word {
			childNode, ok := node.Children.Load(char) // 使用 Load 方法获取子节点
			if !ok {
				// 如果单词不存在于前缀树中，则直接返回
				return
			}
			node = childNode.(*TrieNode) // 类型断言
		}
		if node.IsEnd {
			// 单词存在，标记为非结束节点，即移除单词
			node.IsEnd = false
		}
	}
}

// RemoveAll 移除前缀树中所有单词
func (t *Trie) RemoveAll() {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.Root = &TrieNode{
		Children: sync.Map{}, // 初始化为 sync.Map
		IsEnd:    false,
	}
}

// Replace 将句子中的违禁词替换为指定的字符串
func (t *Trie) Replace(sentence string, replace string) string {
	t.mu.RLock()
	defer t.mu.RUnlock()
	// 使用 strings.Builder 高效构建字符串
	var replacedSentence strings.Builder
	textRunes := []rune(sentence)
	i := 0
	for i < len(textRunes) {
		subText := string(textRunes[i:])
		if word, size := findLongestWordInTrie(subText, t); word != "" {
			// 替换为指定字符串
			replacedSentence.WriteString(strings.Repeat(replace, len([]rune(word))))
			i += size
		} else {
			// 保留原始字符
			replacedSentence.WriteRune(textRunes[i])
			i++
		}
	}
	return replacedSentence.String()
}

// LoadReader 实现了 io.Reader 接口的读取器，例如文件
func (t *Trie) LoadReader(reader io.Reader) error {
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		word := strings.TrimSpace(scanner.Text())
		if word != "" {
			t.Add(word)
		}
	}
	return scanner.Err()
}
